#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
修复用户统计表结构
"""

import pymysql
from database.database import engine
from sqlalchemy import text

def fix_user_stats_table():
    """修复用户统计表结构"""
    print("🔧 开始修复用户统计表结构...")
    
    try:
        # 获取数据库连接
        connection = engine.raw_connection()
        cursor = connection.cursor()
        
        # 检查表是否存在
        cursor.execute("SHOW TABLES LIKE 'user_stats'")
        if not cursor.fetchone():
            print("❌ 用户统计表不存在，需要重新创建")
            return False
        
        # 检查created_at字段是否存在
        cursor.execute("DESCRIBE user_stats")
        columns = [column[0] for column in cursor.fetchall()]
        
        if 'created_at' not in columns:
            print("📝 添加缺失的created_at字段...")
            try:
                cursor.execute("""
                    ALTER TABLE user_stats 
                    ADD COLUMN created_at DATETIME DEFAULT CURRENT_TIMESTAMP
                """)
                connection.commit()
                print("✅ created_at字段添加成功")
            except Exception as e:
                print(f"❌ 添加created_at字段失败: {e}")
                connection.rollback()
                return False
        else:
            print("✅ created_at字段已存在")
        
        # 检查其他必要字段
        required_columns = ['id', 'user_id', 'stat_date', 'tts_count', 'asr_seconds']
        missing_columns = []
        
        for column in required_columns:
            if column not in columns:
                missing_columns.append(column)
        
        if missing_columns:
            print(f"❌ 缺少必要字段: {missing_columns}")
            return False
        
        print("✅ 用户统计表结构检查完成")
        
        # 显示当前表结构
        print("\n📋 当前表结构:")
        cursor.execute("DESCRIBE user_stats")
        for column in cursor.fetchall():
            print(f"  {column[0]}: {column[1]}")
        
        cursor.close()
        connection.close()
        return True
        
    except Exception as e:
        print(f"❌ 修复表结构失败: {e}")
        return False

def recreate_user_stats_table():
    """重新创建用户统计表"""
    print("🔄 重新创建用户统计表...")
    
    try:
        # 获取数据库连接
        connection = engine.raw_connection()
        cursor = connection.cursor()
        
        # 删除旧表（如果存在）
        cursor.execute("DROP TABLE IF EXISTS user_stats")
        print("🗑️ 删除旧表")
        
        # 创建新表
        create_table_sql = """
        CREATE TABLE user_stats (
            id INT AUTO_INCREMENT PRIMARY KEY,
            user_id INT NOT NULL,
            stat_date DATE NOT NULL DEFAULT (CURRENT_DATE),
            tts_count INT DEFAULT 0,
            asr_seconds INT DEFAULT 0,
            created_at DATETIME DEFAULT CURRENT_TIMESTAMP,
            FOREIGN KEY (user_id) REFERENCES users(id),
            UNIQUE KEY unique_user_date (user_id, stat_date)
        ) COMMENT='用户每日TTS/ASR使用统计'
        """
        
        cursor.execute(create_table_sql)
        connection.commit()
        print("✅ 用户统计表创建成功")
        
        # 显示新表结构
        print("\n📋 新表结构:")
        cursor.execute("DESCRIBE user_stats")
        for column in cursor.fetchall():
            print(f"  {column[0]}: {column[1]}")
        
        cursor.close()
        connection.close()
        return True
        
    except Exception as e:
        print(f"❌ 重新创建表失败: {e}")
        return False

def main():
    """主函数"""
    print("🚀 开始修复用户统计表")
    
    # 尝试修复现有表
    if fix_user_stats_table():
        print("✅ 表结构修复成功")
    else:
        print("⚠️ 修复失败，尝试重新创建表...")
        if recreate_user_stats_table():
            print("✅ 表重新创建成功")
        else:
            print("❌ 表重新创建失败")
            return
    
    print("\n🎉 用户统计表修复完成！")

if __name__ == "__main__":
    main() 